import os
import csv
from config import dummy_data_dir

output_file = os.path.join(dummy_data_dir, '家庭宽带办理信息表.csv').replace(os.sep, '/')
print(output_file)


def load_residential_info():
    """
    从addr_info.csv文件中加载住宅信息
    
    Returns:
        list: 包含住宅信息的字典列表，每个元素包含district、name、lat、lng
    """
    residential_info = []
    csv_path = os.path.join(dummy_data_dir, 'addr_info.csv').replace(os.sep, '/')
    
    try:
        with open(csv_path, 'r', encoding='utf-8') as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                if row.get('type', '').strip() == '住宅':
                    residential_info.append({
                        "district": row['district'],
                        "name": row['name'],
                        "lat": float(row['latitude']),
                        "lng": float(row['longitude'])
                    })
        print(f"成功加载 {len(residential_info)} 条住宅信息")
        return residential_info
    except FileNotFoundError:
        print(f"错误：找不到文件 {csv_path}")
        return []
    except Exception as e:
        print(f"读取CSV文件时出错: {e}")
        return []


def generate_brd_csv():
    """
    生成包含user_id和brd_addr字段的CSV文件
    user_id格式为U00001-U03000
    brd_addr由district、name和随机生成的楼栋地址及门牌号组成
    每个地址按指定比例随机对应1-4个user_id：1人=30%, 2人=20%, 3人=40%, 4人=10%
    总数据行数为3000条
    """
    import random
    
    # 加载住宅信息
    local_info = load_residential_info()
    if not local_info:
        print("无法加载住宅信息，程序终止")
        return
    
    # 定义楼栋地址和门牌号的模板
    building_numbers = [f"{i}栋" for i in range(1, 21)]  # 1-20栋
    room_numbers = [f"{i:02d}{j}" for i in range(1, 21) for j in ['1', '2', '3']]  # 011-203
    
    # 准备CSV数据
    csv_data = []
    
    # 生成3000条记录，每个地址随机对应1-4个user_id
    user_id_counter = 1
    max_users = 3000
    
    while user_id_counter <= max_users:
        # 随机选择一个位置
        location = random.choice(local_info)
        district = location["district"]
        name = location["name"]
        base_addr = f"{district}{name}"
        
        # 随机选择楼栋和门牌号
        building = random.choice(building_numbers)
        room = random.choice(room_numbers)
        
        # 组合完整地址
        brd_addr = f"{base_addr}{building}{room}"
        
        # 根据指定比例随机决定这个地址对应的user_id数量
        # 比例：1人=0.3, 2人=0.2, 3人=0.4, 4人=0.1
        rand_val = random.random()
        if rand_val < 0.3:
            users_per_addr = 1
        elif rand_val < 0.5:  # 0.3 + 0.2 = 0.5
            users_per_addr = 2
        elif rand_val < 0.9:  # 0.5 + 0.4 = 0.9
            users_per_addr = 3
        else:  # 0.9 + 0.1 = 1.0
            users_per_addr = 4
        
        # 为当前地址生成user_id
        for j in range(users_per_addr):
            if user_id_counter > max_users:
                break
            user_id = f"U{user_id_counter:05d}"
            csv_data.append([user_id, brd_addr])
            user_id_counter += 1
    
    # 写入CSV文件
    with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['user_id', 'brd_addr'])  # 写入表头
        writer.writerows(csv_data)  # 写入数据
    
    print(f"CSV文件已生成: {output_file}")
    print(f"共生成 {len(csv_data)} 条记录")


# 调用函数生成CSV文件
if __name__ == "__main__":
    generate_brd_csv()